// Copyright (c) 2011 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
//
// Author: Markus Gutschke

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "base64.h"
#include "check.h"
#include "http.h"
#include "multipart.h"
#include "misc.h"
#include "stringprintf.h"

#define min(a,b) ({ typeof(a) _a=(a); typeof(b) _b=(b); _a < _b ? _a : _b; })


struct MultiPartEntry {
  char   *content;
  size_t content_length;
  char   *content_type;
};

static void destroyMultiPartEntry(void *arg, char *key, char *value) {
  struct MultiPartEntry *entry = (struct MultiPartEntry *)value;
  if (entry) {
    free(entry->content);
    free(entry->content_type);
  }
  free(key);
  free(value);
}

struct MultiPart *newMultiPart() {
  struct MultiPart *multipart;
  check(multipart = malloc(sizeof(struct MultiPart)));
  initMultiPart(multipart);
  return multipart;
}

void initMultiPart(struct MultiPart *multipart) {
  initHashMap(&multipart->contents, destroyMultiPartEntry, NULL);
}

void destroyMultiPart(struct MultiPart *multipart) {
  if (multipart) {
    destroyHashMap(&multipart->contents);
  }
}

void deleteMultiPart(struct MultiPart *multipart) {
  destroyMultiPart(multipart);
  free(multipart);
}

void addToMultiPart(struct MultiPart *multipart, const char *name,
                    const char *content, ssize_t content_length,
                    const char *content_type) {
  struct MultiPartEntry *entry;
  entry = (struct MultiPartEntry *)getFromHashMap(&multipart->contents, name);
  if (!entry) {
    check(name = strdup(name));
    check(entry = calloc(1, sizeof(struct MultiPartEntry)));
    addToHashMap(&multipart->contents, name, (char *)entry);
  }
  if (content_length < 0) {
    content_length = strlen(content);
  }
  check(entry->content = realloc(entry->content, content_length));
  memcpy(entry->content, content, content_length);
  entry->content_length = content_length;
  check(entry->content_type = realloc(entry->content_type,
                                      strlen(content_type) + 1));
  strcpy(entry->content_type, content_type);
}

struct MultiPartInfo {
  char   *response;
  size_t *boundaryPos;
  size_t len;
};

static int serializeEntry(void *arg, const char *key, char **value) {
  struct MultiPartInfo *info = (struct MultiPartInfo *)arg;
  struct MultiPartEntry *entry = (struct MultiPartEntry *)*value;

  char *name = httpHeaderEncoding(key);
  info->response = stringPrintf(info->response,
    "\r\n"
    "Content-Type: %s\r\n"
    "Content-Disposition: form-data; name=\"%s\"\r\n",
    entry->content_type ? entry->content_type : "application/octet-stream",
    name);
  free(name);

  char *content_ptr = entry->content;
  for (size_t i = 0, linelen = 0;
       i < entry->content_length;
       ++i, ++content_ptr, ++linelen) {
    if (*content_ptr == '\r' || *content_ptr == '\n') {
      linelen = -1;
    } else if (*content_ptr <= '\000' /* || *content_ptr > '\x7E' ||
               linelen > 75 */) {
      // NOTE: For now, we only switch to BASE64 encoding, if we are absolutely
      // sure we can't handle the contents otherwise. The server has
      // difficulty dealing with Content-Transfer-Encodings.

      // Content is potentially not safe to transmit verbatim. Better encode
      // with Base64
      info->response = stringPrintf(info->response,
        "Content-Transfer-Encoding: base64\r\n\r\n");
      size_t offset = strlen(info->response);
      size_t len = (entry->content_length+2)/3*4;
      size_t total = offset + len + (len+75)/76*2 + UNIQUE_ID_LEN + 3;
      check(info->response = realloc(info->response, total));
      for (size_t i = 0; i < entry->content_length; ) {
        int raw_len = min(entry->content_length-i, 76/4*3);
        b64encode(info->response + offset, entry->content+i, raw_len);
        i += raw_len;
        offset += strlen(info->response + offset);
        strcpy(info->response + offset, "\r\n");
        offset += 2;
      }
      sprintf(info->response + offset, "--%*s", UNIQUE_ID_LEN, "");
      info->boundaryPos[info->len++] = strlen(info->response) - UNIQUE_ID_LEN;
      check(strlen(info->response) == total - 1);
      return 1;
    }
  }

  size_t offset = strlen(info->response);
  check(info->response = realloc(info->response,
                                 offset + 2 + entry->content_length + 5 +
                                 UNIQUE_ID_LEN));
  sprintf((char *)memcpy(strcpy(
    info->response + offset, "\r\n") + 2, entry->content,entry->content_length)
    + entry->content_length, "\r\n--%*s", UNIQUE_ID_LEN, "");
  info->boundaryPos[info->len++] = strlen(info->response) - UNIQUE_ID_LEN;
  return 1;
}

static const char *multiEncodeMultiPart(const struct MultiPart *multipart) {
  // Initial part of the HTTP header
  struct MultiPartInfo info;
  info.response = stringPrintf(NULL,
    "MIME-Version: 1.0\r\n"
    "Content-Type: multipart/form-data; boundary=\"%*s\"\r\n"
    "\r\n"
    "--%*s", UNIQUE_ID_LEN, "", UNIQUE_ID_LEN, "");

  // Fill out the "info" data passed into the iterator
  check(info.boundaryPos = malloc(sizeof(size_t) *
                                  (2 + getHashmapSize(&multipart->contents))));
  info.len = 2;
  info.boundaryPos[0] = strcspn(info.response, "=") + 2;
  info.boundaryPos[1] = strlen(info.response) - UNIQUE_ID_LEN;

  // Iterate over all entries, appending their contents to the "response"
  size_t len = info.boundaryPos[1]-2;
  iterateOverHashMap((struct HashMap *)&multipart->contents, serializeEntry,
                     &info);

  // Append trailing "--"
  info.response = stringPrintf(info.response, "--");
  check(info.len == 2+getHashmapSize(&multipart->contents));

  // Set a unique boundary marker
  const char *uniq;
  while (strstr(info.response, uniq = uniqueID())) {
    free((void *)uniq);
  }
  for (size_t i = 0; i < info.len; ++i) {
    memcpy(info.response + info.boundaryPos[i], uniq, UNIQUE_ID_LEN);
  }
  free((void *)uniq);

  // Compute Content-Length header
  len = info.boundaryPos[info.len-1] + UNIQUE_ID_LEN + 2 - len;
  free(info.boundaryPos);
  char *msg = stringPrintf(NULL,
    "Content-Length: %zd\r\n"
    "%s",
    len, info.response);
  free(info.response);
  return msg;
}

static char *urlEncode(const char *s, size_t l) {
  char *res = NULL;
  size_t sz = 0;
  size_t len = 0;
  for (;;) {
    if (sz < len + 3) {
      check(res = realloc(res, sz += 40));
    }
    unsigned char ch;
    if (!l--) {
      res[len++] = '\000';
      check(res = realloc(res, len));
      return res;
    } else if ((ch = (unsigned char)*s++) == ' ') {
      res[len++] = '+';
    } else if (ch < ' ' || ch == '+' || ch == '&' || ch == '?' ||
               ch == '=' || ch == '%' || ch == '#' || ch > 0x7E) {
      sprintf(res + len, "%%%02X", ch);
      len += strlen(res + len);
    } else {
      res[len++] = (char)ch;
    }
  }
}

static int serializeURL(void *arg, const char *key, char **value) {
  char **response = (char **)arg;
  struct MultiPartEntry *entry = (struct MultiPartEntry *)*value;

  char *k = urlEncode(key, strlen(key));
  char *v = urlEncode(entry->content, entry->content_length);
  *response = stringPrintf(*response, "%s=%s&", k, v);
  free(k);
  free(v);
  return 1;
}

static const char *urlEncodeMultiPart(const struct MultiPart *multipart) {
  char *response;
  check(response = strdup(
    "MIME-Version: 1.0\r\n"
    "Content-Type: application/x-www-form-urlencoded\r\n"
    "\r\n"));
  size_t len = strlen(response);
  iterateOverHashMap((struct HashMap *)&multipart->contents, serializeURL,
                     &response);
  size_t l = strlen(response + len);
  if (response[len + l - 1] == '&') {
    response[len + --l] = '\000';
  }
  char *msg = stringPrintf(NULL,
    "Content-Length: %zd\r\n"
    "%s",
    l, response);
  free(response);
  return msg;
}

const char *serializeMultiPart(const struct MultiPart *multipart,
                               enum MultiPartEncoding encoding) {
  switch (encoding) {
  case M_MULTIPART:
    return multiEncodeMultiPart(multipart);
  case M_URLENCODED:
    return urlEncodeMultiPart(multipart);
  default:
    check(0);
  }
}
